// Copyright 2022 jeff.li. and/or its affiliates.

#include <gtest/gtest.h>
#include <tbir/runtime/container.h>
#include <tbir/runtime/container/iterator_adapter.h>
#include <tbir/runtime/registry.h>
#include <iostream>
#include <random>

namespace tbir::runtime {

    TBIR_REGISTER_GLOBAL("test.PackedFunc_FloatTensorMaker").set_body_typed([](List shape) {
        // automatically convert arguments to desired type.
        using Converter = GenericValueConverter<int64_t>;
        using IterAdapter = IteratorAdapter<Converter, List::iterator>;
        std::vector<int64_t> stl_shape{IterAdapter(shape.begin()), IterAdapter(shape.end())};
        DataType dtype = DataType::Float(32);
        auto result = NDArray::Empty(stl_shape, dtype, {kDLCPU, 0});
        auto ptr_b = static_cast<float *>(result->data);
        for (size_t i = 0; i < 6; ++i) {
            ptr_b[i] = 1.1;
        }
        // automatically assign value return to rv
        return result;
    });

    TEST(PackedFunc, GenericTest) {
        const auto *test_func =
                ::tbir::runtime::FunctionRegistry::Get("test.PackedFunc_FloatTensorMaker");
        DLContext ctx;
        ctx.device_type = kDLCPU;
        ctx.device_id = 0;
        List shape = {2, 3};
        NDArray tsr = (*test_func)({RTView(shape)}).As<NDArray>();
        std::cout << tsr << std::endl;
    }

}  // namespace tbir::runtime
